import os
import shutil
from tqdm import tqdm
import random
import json
import cv2
# import argparse

# def parse_args():
#     """
#     Set args parameters
#     """
#     parser = argparse.ArgumentParser(description='Seg model turnon test.')
#     parser.add_argument("--img_dir", type=str, default='/opt/data/private/ESD_monitor/data/seg/seg3/image', help="Define the data location.")
#     parser.add_argument("--ckpt_dir", type=str, default='/root/cwq/t7-classification-train-master/seg_model/ESD/model.pth', help="Define where to save model checkpoints.")
#     parser.add_argument("--save_dir", type=str, default='./result12', help="Define where to save evaluate results.")
#     args = parser.parse_args()

#     return args

headstr = """\
<annotation>
    <folder>%s</folder>
    <filename>%s</filename>
    <path>auto_repair</path>
    <source>
        <database>Unknown</database>
    </source>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>
"""

objstr = """\
    <object>
        <name>%s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>%d</xmin>
            <ymin>%d</ymin>
            <xmax>%d</xmax>
            <ymax>%d</ymax>
        </bndbox>
    </object>
"""

tailstr = """\
</annotation>
"""
def iou(b1, b2):
    # [x1, y1, x2, y2]
    ltx = max(b1[0], b2[0])
    lty = max(b1[1], b2[1])
    brx = min(b1[2], b2[2])
    bry = min(b1[-1], b2[-1])
    if ltx < brx and lty < bry:
        return True
    return False


def find_one(box):
    allx = [x[0] for x in box]
    ally = [x[1] for x in box]
    return [min(allx), min(ally), max(allx), max(ally)]


def get_bbox_resized(point, h, w):
    all_box = find_one(point)
    xmin = int(max(0, all_box[0]))
    ymin = int(max(0, all_box[1]))
    xmax = int(min(w, max(xmin+5, all_box[2])))
    ymax = int(min(h, max(ymin+5, all_box[3])))
    labelbox = [xmin, ymin, xmax, ymax]

    # 额外选取一个没有缺陷的other类
    r = []

    rx = random.randint(10, w-225)
    ry = random.randint(10, h-225)

    otherbox = [rx, ry, rx + 224, ry + 224]
    if not iou(labelbox, otherbox):
        r = otherbox    
    return [labelbox, r]


def get_bbox_fixed(point, h , w):
    all_box = find_one(point)

    cx = (all_box[2] + all_box[0])//2
    cy = (all_box[3] + all_box[1])//2

    # 如果左上角有问题，设为0，那么右下角不会出问题
    xmin = int(max(cx - 112, 0))
    ymin = int(max(cy - 112, 0))

    xmax = int(min(xmin + 224, w))
    ymax = int(min(ymin + 224, h))

    # 如果右下角有问题，左上角应该没问题
    if xmax == w:
        xmin = xmax - 224
    if ymax == h:
        ymin = ymax - 224


    labelbox = [xmin, ymin, xmax, ymax]

    # 额外选取一个没有缺陷的other类
    r = []

    rx = random.randint(10, w-225)
    ry = random.randint(10, h-225)

    otherbox = [rx, ry, rx + 224, ry + 224]
    if not iou(labelbox, otherbox):
        r = otherbox    
    return [labelbox, r]

def defect_mapping(defect):
    res = defect[0:5]
    if len(defect) >5:
        last_char = defect[-1]
        if '1' == last_char:
            res = 'CUT'
        #elif '3' == last_char:
        #    res = 'EXTEND'
    return res

def write_xml(path, objstr, head, objs, tail):
    # path = os.path.join(new_xml_file, p+".xml")
    # print(objs)
    with open(path, 'w', encoding='utf-8') as f:
        f.write(head)
        for obj in objs:
            f.write(objstr % (obj[0], obj[1], obj[2], obj[3], obj[4]))
        f.write(tail)
    
def main(working_dir, saving_dir):
    for cur, dirs, files in os.walk(working_dir):
        print(cur, dirs, len(files))
        if len(dirs) > 0:
            continue
        all_jsons = [os.path.join(cur, i) for i in files if i.endswith('.json')]
        all_imgs = [os.path.join(cur, i) for i in files if i.endswith('.jpg')]
        

        for img_path in tqdm(all_imgs):
            img = cv2.imread(img_path)
            json_file = img_path.replace('jpg', 'json')
            height, width, _ = img.shape
            jpgname = os.path.basename(img_path)

            if json_file in all_jsons:
                with open(json_file , "r") as f:
                    one_json = json.load(f)
                    labels = one_json['shapes']

                ct = 0
                for label in labels:
                    #res = get_bbox_resized(label['points'], height, width)
                    [res, temp] = get_bbox_fixed(label['points'], height, width)
                    bbox = res
                    xmin = bbox[0]
                    ymin = bbox[1]
                    xmax = bbox[2]
                    ymax = bbox[3]

                    if ymin >= ymax or xmin >= xmax:
                        continue

                    # other = res[1]

                    defect = defect_mapping(label['label'])
                    # if defect == 'CUT':
                        # continue

                    head = headstr % (defect, jpgname, int(width), int(height))
                    
                    # temp = '/opt/data/private/t7-aa-cls/annotations'
                    # saving_json = os.path.join(temp, defect)
                    saving_json = os.path.join(saving_dir, 'annotations', defect)
                    saving_img = os.path.join(saving_dir, 'images', defect)
                    os.makedirs(saving_json, exist_ok=True)
                    os.makedirs(saving_img, exist_ok=True)
                    filename = f"{os.path.basename(img_path)[:-4]}_{ct}"
                    xml_path = os.path.join(saving_json, f"{filename}.xml")
                    output_img = os.path.join(saving_img, f"{filename}.jpg")

                    # objs = [[defect, max(0,int(xmin)), max(0,int(ymin)), min(int(xmax), int(width)), min(int(ymax), int(height))]]
                    objs = [[defect, 0, 0, 224, 224]]
                    write_xml(xml_path, objstr, head, objs, tailstr)

                    patch = img[ymin:ymax, xmin:xmax, ...]
                    new_size = (224, 224)
                    patch_resized = cv2.resize(patch, new_size)
                    cv2.imwrite(output_img, patch_resized)
                    
                    ct += 1

            else:
                # 没有缺陷
                if 'TGXID' not in img_path:
                    print(f"111111{img_path}")
                    continue
                # defect = "TGXID"
                #
                # rx = random.randint(1, width - 225)
                # ry = random.randint(1, height - 225)
                # padding = 224
                #
                # if (rx+ry)%5 == 0:
                #     # randcrop = random.randint(padding + 20, 224)
                #     other = [rx, ry, rx + padding, ry + padding]
                #
                #     xmin2 = other[0]
                #     ymin2 = other[1]
                #     xmax2 = other[2]
                #     ymax2 = other[3]
                #
                #     # temp = '/opt/data/private/t7-aa-cls/annotations'
                #     # saving_json = os.path.join(temp, defect)
                #     saving_json = os.path.join(saving_dir, 'annotations', defect)
                #     saving_img = os.path.join(saving_dir, 'images', defect)
                #     os.makedirs(saving_json, exist_ok=True)
                #     os.makedirs(saving_img, exist_ok=True)
                #     filename = os.path.basename(img_path)[:-4]
                #     xml_path = os.path.join(saving_json, f"{filename}.xml")
                #     output_img = os.path.join(saving_img, f"{filename}.jpg")
                #
                #     head = headstr % (defect, jpgname, int(width), int(height))
                #     # objs = [["others", max(0,int(xmin2)), max(0,int(ymin2)), min(int(xmax2), int(w)), min(int(ymax2), int(h))]]
                #     objs = [[defect, 0, 0, 224, 224]]
                #     write_xml(xml_path, objstr, head, objs, tailstr)
                #
                #     patch = img[ymin2:ymax2, xmin2:xmax2, ...]
                #     new_size = (224, 224)
                #     patch_resized = cv2.resize(patch, new_size)
                #     cv2.imwrite(output_img, patch_resized)


if __name__ == "__main__":
    working_dir = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/all_img_cut_pad'
    saving_dir = r'/root/project/AutoRepair_T7/data/t7_training_data/t7-aa-1218-ori/all_img_cut_pad_for_classification'
    if os.path.exists(saving_dir):
        shutil.rmtree(saving_dir)
    os.makedirs(saving_dir, exist_ok=True)

    main(working_dir, saving_dir)
